#  Copyright 2024-2034 the original author or authors.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#  http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import torch
import torchvision.transforms as T
from PIL import Image, ImageOps
from torchvision.models import resnet50

import numpy as np
from facenet_pytorch import MTCNN, InceptionResnetV1

def recognitionTest() :

    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    # 加载MTCNN模型用于人脸定位
    mtcnn = MTCNN(device=device)

    # 加载InceptionResnetV1模型用于人脸识别
    model = InceptionResnetV1(pretrained='vggface2').eval().to(device)

    # 图片路径
    image_path = 'path_to_image.jpg'

    # 将图片转换为PIL图片
    image = Image.open(image_path).convert('RGB')

    # 使用MTCNN模型定位图片中的所有人脸
    faces, probs, bboxes = mtcnn(image, return_probabilities=True, return_bboxes=True)

    # 检查是否找到了人脸
    if len(faces) == 0:
        raise RuntimeError('Unable to find a face')

    # 如果找到多于一张脸，只选取最大的脸
    if len(faces) > 1:
        sorted_indices = np.argsort([bbox.area() for bbox in bboxes])
        faces = [faces[i] for i in sorted_indices[-1:]]
        bboxes = [bboxes[i] for i in sorted_indices[-1:]]
        probs = [probs[i] for i in sorted_indices[-1:]]

    # 取最大的脸进行人脸识别
    face = faces[0].to(device)

    # 标准化输入图片
    transform = T.Compose([
        T.Resize((299, 299)),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    # 对人脸进行识别
    with torch.no_grad():
        output = model(transform(face).unsqueeze(0))

    # 获取最高分类作为识别结果
    index = output.argmax(dim=1).item()
    print(f'Face recognition result: {model.classes[index]}')

    # 保存识别结果
    result_image = image.copy()
    for bbox in bboxes:
        bbox.draw(result_image, color='#00FF00', width=2)
    result_image.save('result.jpg')

if '__main__' == __name__ :
    recognitionTest()